"""Module responsible for depth image processing"""
from multiprocessing import Process, Queue
from PyQt4 import QtGui, QtCore

import numpy as np
import re
import scipy.io as sio
import scipy.ndimage as sn

import params
import _cam as cam

from time import clock, sleep


class ImageProcessor():
    def __init__(self):
        self.runs = 0 #current iteration
        self.initial = True #initial phase - calibration
        self.handsTracked = False 
        self.time = None
        
    def reinit(self):     
        """Set initial phase (calibration) and resets FPS counter every 100 cycles"""
        if not self.runs:
            self.time = clock()
        self.runs += 1
        if self.runs > params.INITIAL_FRAMES:
            if self.initial == True:
                self.initial = False
            self.runs = 0
            self.time = clock()
                      
        self.handsTracked = False
    
    def processImage(self):
        """Main processing routine"""
        if self.initial:
            #position of tracking window, its format: x left, y top, width, height
            self.initRPos = Position(220, 200, 100, 100)
            self.initLPos = Position(370, 200, 100, 100)
            
            #in initial phase create first template frames from both hands and compute parameters         
            self.lParams = self.computeParams(self.initLPos)
            self.rParams = self.computeParams(self.initRPos)
                      
        else:
            self.lParams = self.traceHand(self.lParams)
            self.rParams = self.traceHand(self.rParams)
            self.handsTracked = True
            
    def traceHand(self, par):
        """Trace single hand using previous hand parameters"""
        #create search window in such way it does not exceed image dimensions
        y1 = par.position.y - params.SEARCH_WINDOW
        if y1 < 0:
            y1 = 0
        y2 = par.position.y + par.position.height + params.SEARCH_WINDOW
        if y2 > 480:
            y2 = 480
        x1 = par.position.x - params.SEARCH_WINDOW
        if x1 < 0 :
            x1 = 0
        x2 = par.position.x + par.position.width + params.SEARCH_WINDOW
        if x2 > 640:
            x2 = 640
        searchWindow = self.depthImage[y1:y2][:, x1:x2]
                                       
        #find objects with similar z position and area as previous frame      
        searchMap = (searchWindow < par.avgDistance + 5)
        hMap, numObj = sn.label(searchMap, np.array([[ 1, 1, 1], [ 1, 1, 1], [ 1, 1, 1]], dtype = bool))
        
        if numObj:
            objsArea = sn.sum(searchMap, hMap, range(numObj + 1))
            objsArea = [(objsArea[i], i) for i in range(numObj + 1)]
            objsArea.sort(key = lambda id: id[0] , reverse = True)
            
            cand = []
            for i in objsArea:
                if 3 * par.area > i[0] > 0.3 * par.area and i[0] > 2000:
                    cand.append(i)        
            
#            self.numSimilar = len(cand)
#            self.hMap = hMap
#            self.numObj = numObj
            if len(cand):             
                objPos = re.findall('\d+', str(sn.find_objects(hMap, cand[-1][1] + 1)))
            
                newPos = []  
                for i in range(len(cand)):
                    width = int(objPos[4 * i + 3]) - int(objPos[4 * i + 2])
                    width = width if width < (11000 / par.avgDistance) else int(11000 / par.avgDistance)
                    height = int(objPos[4 * i + 1]) - int(objPos[4 * i])
                    height = height if height < (11000 / par.avgDistance) else int(11000 / par.avgDistance)
                    newPos.append(Position(x1 + int(objPos[4 * i + 2]), y1 + int(objPos[4 * i]),
                                           width, height))      
               
                coeffs = []
                for i in newPos:
                    coeffs.append(i / par.position)
                
                if np.min(coeffs) < 60:                    
                    par = self.computeParams(newPos[np.argmin(coeffs)])
 

#                self.searchMap = searchMap
        
        return par
        
    def computeParams(self, pos):
        """Get parameters of hand in provided image, function assumes the biggest object is hand"""
        im = self.depthImage[pos.y:pos.y + pos.height][:, pos.x:pos.x + pos.width]
        hPar = HandParams()
        
#        y, x = im.shape
        imMap = (im < im.min() + 7)

            
        hMap, numObj = sn.label(imMap, np.array([[ 1, 1, 1], [ 1, 1, 1], [ 1, 1, 1]], dtype = bool))
        
        objsArea = sn.sum(imMap, hMap, range(numObj + 1))
        objsArea = [(objsArea[i], i) for i in range(numObj + 1)]
        objsArea.sort(key = lambda id: id[0] , reverse = True)
            
        hId = objsArea[0][1]
        hPar.area = objsArea[0][0]
        
        objPos = re.findall('\d+', str(sn.find_objects(hMap)[hId - 1]))
        hPar.position = Position(pos.x + int(objPos[2]), pos.y + int(objPos[0]), int(objPos[3]) - int(objPos[2]),
                                 int(objPos[1]) - int(objPos[0]))
        
        hPar.silhouette = (hMap == hId)
        hPar.silDepth = np.where(hMap == hId, im, 0)
        hPar.avgDistance = hPar.silDepth.sum() / hPar.area
        
        tmp = sn.center_of_mass(hPar.silhouette)
        hPar.centerOfWeight = Point(pos.x + int(tmp[1]), pos.y + int(tmp[0]))
        return hPar

    def getDepthImage(self):   
        """Returns depth image with first plane objects only as numpy array""" 
        self.depthImage = cam.getDepthImageRaw() 
        b = 1085 - 32573 / params.MAX_DISTANCE
        np.clip(self.depthImage, 0, b, self.depthImage)
        self.depthImage = 100 / (-0.00307 * self.depthImage + 3.33)
        #weave.blitz("depthImage = 100 / ( -0.00307 * depthImage + 3.33 )") // if weave will be python3 compatible      
        return self.depthImage
      
    def getDepthImageWithSelection(self):
        """Returns depth image converted from array to QPixmap object"""    
        maxValue = np.max(self.depthImage) 
        minValue = np.min(self.depthImage)
        depthIm = np.array((self.depthImage - minValue) / maxValue * 255, np.uint8)     
        img = depth2QImage(depthIm)
        pixmap = QtGui.QPixmap(img)

        pen = QtGui.QPen()
        draw = QtGui.QPainter(pixmap)
        
        if self.initial:
            #draw hud          
            draw.setRenderHint(QtGui.QPainter.Antialiasing, True)
           
            pen.setColor(QtGui.QColor(92, 178, 244, 255))
            draw.setPen(pen)
            
            #draw text
            font = QtGui.QFont("Segoe UI", 10, 75)
            draw.setFont(font)
            draw.drawText(200, 320, "Place your hands and wait for rectangles to disappear")
            draw.drawText(240, 190, "Right hand")
            draw.drawText(390, 190, "Left hand")
            draw.drawText(200, 340, "Ticks left: " + str(params.INITIAL_FRAMES - self.runs))
            
            #draw rectangles
            pen.setWidth(2)
            pen.setColor(QtGui.QColor(255, 255, 255, 255))
            draw.setPen(pen)
            draw.setBrush(QtGui.QBrush(QtGui.QColor(255, 233, 113, 50)))
            draw.drawRect(220, 200, 100, 100)            
            draw.drawRect(370, 200, 100, 100)
            
            #draw circles      
            pen.setWidth(0)
            pen.setColor(QtGui.QColor(255, 255, 255, 0))
            draw.setPen(pen)
            draw.setBrush(QtGui.QBrush(QtGui.QColor(255, 255, 255, 128)))
            center = QtCore.QPoint(270, 250)
            draw.drawEllipse(center, 6, 6)
            center = QtCore.QPoint(420, 250)
            draw.drawEllipse(center, 6, 6)
        else:
            #mark selected hands
            pen.setWidth(2)
            pen.setColor(QtGui.QColor(255, 255, 255, 255))
            draw.setPen(pen)
            draw.setBrush(QtGui.QBrush(QtGui.QColor(255, 233, 113, 50))) 
            draw.drawRect(*self.lParams.position())
            draw.drawRect(*self.rParams.position())   
            
            #draw circles      
            pen.setWidth(0)
            pen.setColor(QtGui.QColor(255, 255, 255, 0))
            draw.setPen(pen)
            draw.setBrush(QtGui.QBrush(QtGui.QColor(255, 255, 255, 128)))
            center = QtCore.QPoint(*self.lParams.centerOfWeight())
            draw.drawEllipse(center, 6, 6)
            center = QtCore.QPoint(*self.rParams.centerOfWeight())
            draw.drawEllipse(center, 6, 6)
            
#            #mark candidates for new objects
#            draw.setPen(QtGui.QPen(QtGui.QColor(92, 178, 244, 255)))
#    
#            font = QtGui.QFont("Segoe UI", 16, 75)
#            draw.setFont(font)
#            draw.drawText(200, 340, "NUMOBJECTS: " + str(self.numObj))
#            draw.drawText(220, 360, "NUMSIMILAR: " + str(self.numSimilar))
#            
#            kolor = []
#            kolor.append(None)
#            kolor.append(QtGui.QColor(255, 0, 0, 255))
#            kolor.append(QtGui.QColor(0, 255, 0, 255))
#            kolor.append(QtGui.QColor(0, 0, 255, 255))
#            kolor.append(QtGui.QColor(255, 255, 0, 255))
#            kolor.append(QtGui.QColor(255, 0, 255, 255))
#            kolor.append(QtGui.QColor(0, 255, 255, 255))
#            kolor.append(QtGui.QColor(255, 128, 128, 255))
#            kolor.append(QtGui.QColor(128, 255, 128, 255))
#            kolor.append(QtGui.QColor(128, 128, 255, 255))
#            
#            if self.numObj > 9:
#                self.numObj = 9
#            for i in range(1, self.numObj + 1):
#                s = np.argwhere(self.hMap == i)
#                draw.setPen(QtGui.QPen(kolor[i]))
#                for k in s:
#                    draw.drawPoint(k[1] + self.lParams.position.x - params.SEARCH_WINDOW, k[0] + self.lParams.position.y - params.SEARCH_WINDOW)
#            try:
#                s = np.argwhere(self.searchMap == True)
#                draw.setPen(QtGui.QPen(QtGui.QColor(0, 255, 0, 255)))
#                for k in s:
#                    draw.drawPoint(k[1] + +self.lParams.position.x - params.SEARCH_WINDOW,
#                                   k[0] + self.lParams.position.y - params.SEARCH_WINDOW)
#            except Exception:
#                pass
                
                        
        #draw FPS
        font = QtGui.QFont("Segoe UI", 10, 75)
        draw.setFont(font)
        draw.setPen(QtGui.QPen(QtGui.QColor(255, 20, 20, 255)))
        draw.drawText(570, 470, '{0} {1:.1f} '.format("FPS:", self.runs / (clock() - self.time)))
        
        return pixmap
    
    def getRGBImage(self):
        """Returns colour RGB image in QImage format from array"""
        self.RGBImage = cam.getRGBImageRaw()  
        img = rgb2QImage(self.RGBImage) 
        return img
     
class HandParams():
    def __init__(self, centerOfWeight = (None, None), position = None, winCorner = (None, None), area = None, avgDistance = None,
                 silhouette = None, silDepth = None):
        self.centerOfWeight = Point(*centerOfWeight)
        self.position = Position(position)
        self.winCorner = Point(winCorner)
        self.area = area
        self.avgDistance = avgDistance
        self.silhouette = silhouette
        self.silDepth = silDepth
        
class Position():
    def __init__(self, x = None, y = None, width = None, height = None):
        self.x = x
        self.y = y
        self.width = width
        self.height = height    
    def __str__(self):
        return "x:" + str(self.x) + " y:" + str(self.y) + " width:" + str(self.width) + " height:" + str(self.height)
    def __call__(self):
        return self.x, self.y, self.width, self.height
    def __truediv__(self, other):
        return np.sqrt((self.x - other.x) ** 2 + (self.y - other.y) ** 2)

class Point():    
    def __init__(self, x = None, y = None):
        self.x = x
        self.y = y   
    def __str__(self):
        return "x:" + str(self.x) + " y:" + str(self.y)
    def __call__(self):
        return self.x, self.y
    
def rgb2QImage(rgb):
    """Convert the 3D numpy array `rgb` into a 32-bit QImage.  `rgb` must
    have three dimensions with the vertical, horizontal and RGB image axes."""   
    h, w, channels = rgb.shape
    
    # Qt expects 32bit BGRA data for color images:
    bgra = np.empty((h, w, 4), np.uint8, 'C')
    bgra[..., 0] = rgb[..., 2]
    bgra[..., 1] = rgb[..., 1]
    bgra[..., 2] = rgb[..., 0]
    bgra[..., 3].fill(255)
    fmt = QtGui.QImage.Format_RGB32

    result = QtGui.QImage(bgra.data, w, h, fmt)
    result.ndarray = bgra
    return result

def depth2QImage(depth):
    """Convert the depth numpy array into a 32-bit QImage."""  
    h, w = depth.shape

    # Qt expects 32bit BGRA data for color images:
    bgra = np.empty((h, w, 4), np.uint8, 'C')
    bgra[..., 0] = depth
    bgra[..., 1] = depth
    bgra[..., 2] = depth
    bgra[..., 3].fill(255)
    fmt = QtGui.QImage.Format_RGB32

    result = QtGui.QImage(bgra.data, w, h, fmt)
    result.ndarray = bgra
    return result

#def getBDS(image):
#    """Compute Bi-Dimensional-Sum aka sum of values across two axis, normalized to sum up to 1.
#    This replaces histogram and gives better result for depth image"""   
#    bds = np.hstack((image.sum(axis = 0), image.sum(axis = 1)))
#    bds = bds / bds.sum()   
#    return bds
#
#def computeBhCoeff(model, candidate):
#    """Compute Bhattacharyya Coefficient of two discrete distributions
#    result equal to 1 means same pdf function
#    result equal to 0 means no similarity"""
#    return np.sqrt(model * candidate).sum()
